package com.spring.ai.tutorial.rag.evaluation.controller;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.evaluation.RelevancyEvaluator;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.evaluation.EvaluationRequest;
import org.springframework.ai.evaluation.EvaluationResponse;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author yingzi
 * @since 2025/6/11
 */
@RestController
@RequestMapping("/rag/evalutaion")
public class RagEvaluationController {

    private static final Logger logger = LoggerFactory.getLogger(RagEvaluationController.class);
    private final SimpleVectorStore simpleVectorStore;
    private final DashScopeChatModel qwenMaxChatModel;
    private final DashScopeChatModel qwenPlusChatModel;

    public RagEvaluationController(EmbeddingModel embeddingModel, @Qualifier("qwen-max")DashScopeChatModel qwenMaxChatModel,
                                   @Qualifier("qwen-plus")DashScopeChatModel qwenPlusChatModel) {
        this.simpleVectorStore = SimpleVectorStore
                .builder(embeddingModel).build();
        this.qwenMaxChatModel = qwenMaxChatModel;
        this.qwenPlusChatModel = qwenPlusChatModel;
    }

    @GetMapping("/add")
    public void add() {
        logger.info("start add data");
        HashMap<String, Object> map = new HashMap<>();
        map.put("year", 2025);
        map.put("name", "yingzi");
        List<Document> documents = List.of(
                new Document("你的姓名是影子，湖南邵阳人，25年硕士毕业于北京科技大学，曾先后在百度、理想、快手实习，曾发表过一篇自然语言处理的sci，现在是一名AI研发工程师"),
                new Document("你的姓名是影子，专业领域包含的数学、前后端、大数据、自然语言处理", Map.of("year", 2024)),
                new Document("你姓名是影子，爱好是发呆、思考、运动", map));
        simpleVectorStore.add(documents);
    }

    @GetMapping("/evalute")
    public String evalute(@RequestParam(value = "query", defaultValue = "你好，请告诉我影子这个人的身份信息") String query) {
        logger.info("start evalute");
        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .vectorStore(simpleVectorStore)
                        .build())
                .build();

        ChatResponse chatResponse = ChatClient.builder(qwenMaxChatModel)
                .build().prompt(query).advisors(retrievalAugmentationAdvisor).call().chatResponse();

        EvaluationRequest evaluationRequest = new EvaluationRequest(
                // The original user question
                query,
                // The retrieved context from the RAG flow
                chatResponse.getMetadata().get(RetrievalAugmentationAdvisor.DOCUMENT_CONTEXT),
                // The AI model's response
                chatResponse.getResult().getOutput().getText()
        );
        logger.info("evalute request: {}", evaluationRequest);

        RelevancyEvaluator evaluator = new RelevancyEvaluator(ChatClient.builder(qwenPlusChatModel));
        EvaluationResponse evaluationResponse = evaluator.evaluate(evaluationRequest);
        boolean pass = evaluationResponse.isPass();
        logger.info("evalute result: {}", pass);
        return chatResponse.getResult().getOutput().getText();
    }
}
